import wandb
import random
import torch
import numpy as np

from torchvision.utils import make_grid
from lightning import Callback
from io import BytesIO
from PIL import Image
from torchmetrics import Metric
from omegaconf import DictConfig


class LogImagesCallback(Callback):
    def __init__(self) -> None:
        super().__init__()

    def on_validation_epoch_end(self, trainer, pl_module) -> None:
        batch = next(iter(trainer.val_dataloaders))

        _, img_view1, img_view2, _ = batch

        time_indices = [0, img_view1.size(2) // 2, img_view1.size(2) - 1]

        imgs_view1 = [img_view1[0, :, t].cpu() for t in time_indices]
        imgs_view2 = [img_view2[0, :, t].cpu() for t in time_indices]

        combined_imgs = imgs_view1 + imgs_view2
        combined_grid = make_grid(combined_imgs, nrow=3)
        
        trainer.logger.experiment.log({
            "val_images": [wandb.Image(combined_grid, caption="Top: View 1 - first, middle, last timesteps; Bottom: View 2 - first, middle, last timesteps")]
        })

class PrintCallback(Callback):
    def on_train_start(self, trainer, pl_module):
        print("Training is started!")
    def on_train_end(self, trainer, pl_module):
        print("Training is done.")


class LogVolumeGifsCallback(Callback):
    def __init__(self, num_samples: int = 4, log_every_n_epochs:int = 50) -> None:
        super().__init__()
        
        self.num_samples = num_samples
        self.batches = []
        self.first_train_batch = None
        self.log_every_n_epochs = log_every_n_epochs

    def create_gif(self, images, duration=500):
        """
        duration time for each frame in miliseconds
        """
        images = [(np.clip(img, 0, 1) * 255).astype('uint8') for img in images]
        pil_images = [Image.fromarray(img) for img in images]
        
        gif_buffer = BytesIO()
        pil_images[0].save(gif_buffer, format='GIF', save_all=True, append_images=pil_images[1:], loop=0, duration=duration)
        gif_buffer.seek(0)
        return gif_buffer
    
    def on_validation_batch_end(self, trainer, pl_module, outputs, batch, batch_idx, dataloader_idx=0):
        self.batches.append({
            'original': outputs['pixel_values'].detach().cpu(),
            'masked': outputs['masked'].detach().cpu(),
            'reconstructed_volume': outputs['reconstructed_volume'].detach().cpu(),
            'im_paste': outputs['im_paste'].detach().cpu(),
            'batch_idx': batch_idx
        })

    def on_validation_epoch_end(self, trainer, pl_module) -> None:
        if trainer.current_epoch % self.log_every_n_epochs != 0:
            return

        if not self.batches:
            return

        selected_samples = []

        random.shuffle(self.batches)
        selected_batches = random.sample(self.batches, min(self.num_samples, len(self.batches)))

        for batch in selected_batches:
            batch_idx = batch['batch_idx']
            num_samples_in_batch = batch['original'].size(0)
            sample_idx = random.choice(range(num_samples_in_batch))

            selected_samples.append({
                'batch_idx': batch_idx,
                'sample_idx': sample_idx,
                'original': batch['original'][sample_idx].squeeze().to(torch.float32).numpy(),
                'masked': batch['masked'][sample_idx].squeeze().to(torch.float32).numpy(),
                'reconstructed': batch['reconstructed_volume'][sample_idx].squeeze().to(torch.float32).numpy(),
                'im_paste': batch['im_paste'][sample_idx].squeeze().to(torch.float32).numpy(),
            })

        selected_samples = selected_samples[:self.num_samples]

        for kind in ['original', 'masked', 'reconstructed', 'im_paste']:
            gifs = [
                wandb.Video(self.create_gif(sample[kind]), fps=10, format="gif", 
                            caption=f"Batch {sample['batch_idx']} Sample {sample['sample_idx']}")
                for sample in selected_samples
            ]
            trainer.logger.experiment.log({f'Val/{kind.capitalize()}': gifs})

        self.batches.clear()
    

class ImageLoggerMetricPyTorch(Metric):
    def __init__(self, cfg: DictConfig, image_encoder, num_samples: int = 1):
        super().__init__()

        self.cfg = cfg
        self.image_encoder = image_encoder
        self.num_samples = num_samples
        self.batch = None

    def update(self, outputs, batch):
        logits = outputs.logits
        mask = outputs.mask
        mask = mask.unsqueeze(-1).repeat(1, 1, self.image_encoder.config.patch_size * self.image_encoder.config.patch_size * self.image_encoder.config.tubelet_size)
        mask = self.image_encoder.unpatchify(mask)
        reconstructed = self.image_encoder.unpatchify(logits)
        masked_pixel_values = batch * (1 - mask)
        im_paste = batch * (1 - mask) + reconstructed * mask

        self.batch = {
            'original': batch.detach().cpu().to(torch.float32),
            'masked': masked_pixel_values.detach().cpu().to(torch.float32),
            'reconstructed_volume': reconstructed.detach().cpu().to(torch.float32),
            'im_paste': im_paste.detach().cpu().to(torch.float32),
        }

    def compute(self):
        selected_index = random.randint(0, self.batch['original'].size(0) - 1)

        selected_batch = {kind: self.batch[kind][selected_index] for kind in self.batch.keys()}

        return selected_batch

    def log_plots(self):
        selected_batch = self.compute()

        for kind in self.batch.keys():

            img_sequence = selected_batch[kind].squeeze().numpy()  # Shape should be (T, 84, 84)

            gif = wandb.Video(self.create_gif(img_sequence), fps=10, format="gif")
            wandb.log({f'val_images/{kind.capitalize()}': gif})
        
        self.reset()

    def create_gif(self, images, duration=500):
        images = [(np.clip(img, 0, 1) * 255).astype('uint8') for img in images]
        pil_images = [Image.fromarray(img) for img in images]
        
        gif_buffer = BytesIO()
        pil_images[0].save(gif_buffer, format='GIF', save_all=True, append_images=pil_images[1:], loop=0, duration=duration)
        gif_buffer.seek(0)
        return gif_buffer

    def reset(self):
        self.batch = None
